from langchain_core.tools import tool
from typing import Annotated, Optional, List
import logging
from tavily import TavilyClient
import os

logger = logging.getLogger(__name__)

tavily_client = TavilyClient(api_key=os.getenv("TAVILY_API_KEY"))

@tool
def tavily_search(
    query: Annotated[str, "搜索查询"],
    max_results: Annotated[int, "最大结果数"] = 5
) -> dict:
    """使用Tavily搜索引擎搜索信息
    
    执行网络搜索并返回相关结果。
    """
    try:
        response = tavily_client.search(query=query, max_results=max_results)
        logger.info(f"[tavily_search] 搜索: {query}")
        return {
            "success": True,
            "query": query,
            "results": response.get("results", [])
        }
    except Exception as e:
        logger.error(f"[tavily_search] 搜索失败: {e}")
        return {"success": False, "error": str(e)}

TavilySearchTool = [tavily_search]
